package experiment4;

import java.io.*;
import java.util.*;
import java.nio.ByteBuffer;

// 定义常量
class Constants {
    static final int NX = 784; // 输入层节点数
    static final int NB = 256; // 隐藏层节点数
    static final int NY = 10;  // 输出层节点数
    static final double ETA = 0.1; // 学习率
    static final int TRAIN_SIZE = 60000; // 训练集大小
    static final int TEST_SIZE = 10000; // 测试集大小
}

// 定义节点结构
class Node {
    double val; // 节点值
    double bias; // 阈值
    List<Double> weight; // 权值向量

    Node() {
        this.weight = new ArrayList<>();
    }
}

public class NeuralNetwork {
    // 输入、隐藏、输出层的节点
    static Node[] x = new Node[Constants.NX]; // 输入层节点
    static Node[] b = new Node[Constants.NB]; // 隐藏层节点
    static Node[] y = new Node[Constants.NY]; // 输出层节点

    // 误差信号
    static double[] g = new double[Constants.NY]; // 输出层误差信号
    static double[] e = new double[Constants.NB]; // 隐藏层误差信号

    // 训练和测试数据集
    static List<List<Double>> trainx = new ArrayList<>();
    static List<List<Double>> trainy = new ArrayList<>();
    static List<List<Double>> testx = new ArrayList<>();
    static List<List<Double>> testy = new ArrayList<>();

    // 文件路径
    static String imageFilePath = "D:\\Java\\projects\\ti\\AIGC\\src\\experiment4\\MNIST\\train-images.idx3-ubyte";  // 图像数据文件路径
    static String labelFilePath = "D:\\Java\\projects\\ti\\AIGC\\src\\experiment4\\MNIST\\train-labels.idx1-ubyte";  // 标签数据文件路径

    public static void main(String[] args) {
        // 初始化网络
        init();
        // 加载数据
        loadData();
        // 训练网络
        train();
        // 测试网络
        test();
    }

    // 激活函数 sigmoid
    public static double sigmoid(double x) {
        return 1.0 / (1.0 + Math.exp(-x));
    }

    // 生成随机权值
    public static double getRandWeight() {
        return Math.random() * 2 - 1; // (-1, 1)
    }

    // 生成随机阈值
    public static double getRandBias() {
        return Math.random() * 0.02 - 0.01; // (-0.01, 0.01)
    }

    // 初始化网络
    public static void init() {
        for (int i = 0; i < Constants.NX; i++) {
            x[i] = new Node();
            for (int j = 0; j < Constants.NB; j++) {
                x[i].weight.add(getRandWeight()); // 输入层到隐藏层的权值
            }
        }

        for (int i = 0; i < Constants.NB; i++) {
            b[i] = new Node();
            b[i].bias = getRandBias(); // 隐藏层的阈值
            for (int j = 0; j < Constants.NY; j++) {
                b[i].weight.add(getRandWeight()); // 隐藏层到输出层的权值
            }
        }

        for (int i = 0; i < Constants.NY; i++) {
            y[i] = new Node();
            y[i].bias = getRandBias(); // 输出层的阈值
        }
    }

    // 前向传播
    public static void forward(int k) {
        for (int i = 0; i < Constants.NX; i++) {
            x[i].val = trainx.get(k).get(i); // 输入层读取数据
        }

        for (int i = 0; i < Constants.NB; i++) {
            b[i].val = 0; // 清空隐藏层的值
            for (int j = 0; j < Constants.NX; j++) {
                b[i].val += x[j].val * x[j].weight.get(i); // 输入层到隐藏层的加权和
            }
            b[i].val = sigmoid(b[i].val - b[i].bias); // 隐藏层的激活值
        }

        for (int i = 0; i < Constants.NY; i++) {
            y[i].val = 0; // 清空输出层的值
            for (int j = 0; j < Constants.NB; j++) {
                y[i].val += b[j].val * b[j].weight.get(i); // 隐藏层到输出层的加权和
            }
            y[i].val = sigmoid(y[i].val - y[i].bias); // 输出层的激活值
        }
    }

    // 反向传播
    public static void back(int k) {
        for (int i = 0; i < Constants.NY; i++) {
            g[i] = y[i].val * (1 - y[i].val) * (trainy.get(k).get(i) - y[i].val); // 计算输出层误差信号
        }

        for (int i = 0; i < Constants.NB; i++) {
            double res = 0;
            for (int j = 0; j < Constants.NY; j++) {
                res += b[i].weight.get(j) * g[j]; // 计算隐藏层误差信号的中间项
            }
            e[i] = b[i].val * (1 - b[i].val) * res; // 计算隐藏层误差信号
        }

        for (int i = 0; i < Constants.NB; i++) {
            for (int j = 0; j < Constants.NY; j++) {
                b[i].weight.set(j, b[i].weight.get(j) + Constants.ETA * b[i].val * g[j]); // 更新隐藏层到输出层的权值
            }
        }

        for (int i = 0; i < Constants.NY; i++) {
            y[i].bias -= Constants.ETA * g[i]; // 更新输出层的阈值
        }

        for (int i = 0; i < Constants.NX; i++) {
            for (int j = 0; j < Constants.NB; j++) {
                x[i].weight.set(j, x[i].weight.get(j) + Constants.ETA * x[i].val * e[j]); // 更新输入层到隐藏层的权值
            }
        }

        for (int i = 0; i < Constants.NB; i++) {
            b[i].bias -= Constants.ETA * e[i]; // 更新隐藏层的阈值
        }
    }

    // 训练网络
    public static void train() {
        int epoch = 10; // 训练轮数
        int n = trainx.size(); // 训练样本数
        for (int i = 0; i < epoch; i++) {
            System.out.println("Epoch " + (i + 1));
            for (int j = 0; j < n; j++) {
                forward(j); // 前向传播
                back(j);    // 反向传播
            }
        }
    }

    // 测试网络
    public static void test() {
        int n = testx.size(); // 测试样本数
        int correct = 0; // 正确识别数
        for (int i = 0; i < n; i++) {
            forward(i); // 前向传播
            int pred = -1; // 预测结果
            double maxVal = -1; // 最大输出值
            for (int j = 0; j < Constants.NY; j++) {
                if (y[j].val > maxVal) {
                    maxVal = y[j].val;
                    pred = j;
                }
            }
            int real = -1; // 真实结果
            for (int j = 0; j < Constants.NY; j++) {
                if (testy.get(i).get(j) == 1.0) {
                    real = j;
                    break;
                }
            }
            if (pred == real) correct++; // 如果预测正确，计数加一
        }
        System.out.println("Accuracy: " + (correct * 100.0 / n) + "%");
    }

    // 读取MNIST图像数据
    public static List<List<Double>> readMNISTImages(String filePath) {
        List<List<Double>> images = new ArrayList<>();
        try (BufferedInputStream inputStream = new BufferedInputStream(new FileInputStream(filePath))) {
            byte[] header = new byte[16];
            inputStream.read(header); // 读取文件头部信息
            int numImages = ByteBuffer.wrap(header, 4, 4).getInt();
            int numRows = ByteBuffer.wrap(header, 8, 4).getInt();
            int numCols = ByteBuffer.wrap(header, 12, 4).getInt();

            for (int i = 0; i < numImages; i++) {
                List<Double> image = new ArrayList<>();
                for (int j = 0; j < numRows * numCols; j++) {
                    image.add((double) (inputStream.read()) / 255.0); // 将像素值归一化
                }
                images.add(image);
            }
        } catch (IOException e) {
            e.printStackTrace();
        }
        return images;
    }

    // 读取MNIST标签数据
    public static List<List<Double>> readMNISTLabels(String filePath) {
        List<List<Double>> labels = new ArrayList<>();
        try (BufferedInputStream inputStream = new BufferedInputStream(new FileInputStream(filePath))) {
            byte[] header = new byte[8];
            inputStream.read(header); // 读取文件头部信息
            int numLabels = ByteBuffer.wrap(header, 4, 4).getInt();
            for (int i = 0; i < numLabels; i++) {
                List<Double> label = new ArrayList<>(Collections.nCopies(Constants.NY, 0.0));
                label.set(inputStream.read(), 1.0); // 将标签转换为one-hot编码
                labels.add(label);
            }
        } catch (IOException e) {
            e.printStackTrace();
        }
        return labels;
    }

    // 加载数据
    public static void loadData() {
        trainx = readMNISTImages(imageFilePath);
        trainy = readMNISTLabels(labelFilePath);
        testx = readMNISTImages(imageFilePath); // 假设测试集和训练集在同一文件夹下
        testy = readMNISTLabels(labelFilePath);
    }
}
